# -*- coding: UTF-8 -*-
from PIL import Image
import numpy as np
import os
import sys
import json
import math
import re
from enum import Enum

COLOR_NUM_MAX = 20      # 图片的最大颜色数量
H_MAX = 100             # 图片的最大高度
W_MAX = 100             # 图片的最大宽度
CHAR_NUM_MAX = 5000     # 高级弹幕最大字符数限制

# 行尾的空格处理模式：
#   KEEP_SPACE, 保留行尾空格不做处理
#   RM_SPACE, 去除行尾空格(但会保留第一行和最后一行的行尾空格,便于使用锚点记性对齐等操作)
class EndMode(Enum):
    KEEP_SPACE = 0
    RM_SPACE = 1


# 行尾的空格处理模式：
end_mode = EndMode.RM_SPACE

# 符号设置, 注意前景与背景符号的宽度需保持一致，否则生成的图片会错位
fg_char = '█'           # 每个颜色图层前景用符号
bg_char = '  '          # 每个颜色图层背景用符号


# 分割图片相关参数
# 分割模式：
# VERTICAL, 垂直分割
# HORIZONTAL, 水平分割
class SplitMode(Enum):
    VERTICAL = 0
    HORIZONTAL = 1


# 图片分割模式
split_mode = SplitMode.VERTICAL

split_num = 0

# 保存生成结果的统计信息的文件
global info_file

# 配置文件
config_path = "./config.ini"


def init(in_dir, out_dir='./out'):
    """
    初始化
    :param in_dir: 要处理的图片所在目录
    :param out_dir: 处理结果输出目录
    :return: 无
    """
    print('---------------------')
    print("初始化...")

    if not os.path.exists(in_dir):
        print(f'找不到输入图片目录: "{in_dir}"\n请将图片复制到目录: "{in_dir}"')
        input("按回车键退出")
        sys.exit(-1)

    if not os.path.exists(out_dir):
        print(f'新建输出结果目录："{out_dir}"')
        os.makedirs(out_dir)

    # 在输出目录下为输入图片的每张图片建立一个同名的文件夹, 用于存放处理结果
    print('新建文件夹...')

    for entry in os.scandir(in_dir):
        file_name = entry.name.rsplit(".", 1)[0]            # 去除后缀名
        dir_to_create = out_dir + "/" + file_name
        if os.path.exists(dir_to_create):
            continue
        print(f'{dir_to_create}')
        os.makedirs(dir_to_create)

    print("初始化完毕")
    print('---------------------')


def write_txt(char_arr, path):
    """
    将生成的字符矩阵写入文本文件
    :param char_arr: 生成的字符矩阵
    :param path: 要写入的文本文件
    :return: 统计该文件中的字符总数
    """
    height, width = char_arr.shape

    line_no = 0         # 行号
    char_sum = 0        # 统计字符数
    with open(path, "w", encoding='UTF-8') as f:
        for i in char_arr:
            line_no += 1
            line = ''.join(i)
            if end_mode == EndMode.RM_SPACE:
                # 保留第一行和最后一行的行尾空格, 便于使用锚点记性对齐等操作
                if line_no != 1 and line_no != height:
                    line = line.rstrip()
            line += '\n'
            f.write(line)
            char_sum += len(line)
    return char_sum


def gen_txt(arr_shape, color, mask, out_dir):
    """
    根据背景mask生成字符矩阵, 并保存为文本文件
    :param arr_shape: 字符矩阵大小
    :param color: 当前处理的颜色
    :param mask: 背景mask
    :param out_dir: 输出路径
    :return: 无
    """
    global end_mode, fg_char, bg_char, split_mode, split_num, info_file

    # 字符矩阵, 与图片的大小一致, item对应的颜色像素位置设为指定字符, 其它位置设为空格
    height, width = arr_shape[:2]
    char_arr = np.zeros((height, width)).astype(dtype=np.str)
    char_arr[:] = fg_char
    char_arr[mask] = bg_char

    # 如果有alpha通道并且不为FF(完全不透明)则输出的颜色字段包含alpha通道的值
    if len(color) == 6:
        line = color + f'【 R: {int(color[:2], 16)}, G: {int(color[2:4], 16)}, B: {int(color[4:], 16)}】:\n'
    elif len(color) == 8:
        alpha = int(color[6:], 16)
        if alpha == 255:
            color = color[:6]
            line = color + f'【 R: {int(color[:2], 16)}, G: {int(color[2:4], 16)}, B: {int(color[4:], 16)}】:\n'
        else:
            line = color + f'【 R: {int(color[:2], 16)}, G: {int(color[2:4], 16)}, B: {int(color[4:6], 16)}, ' \
                           f'A: {alpha}】:\n'

    info_file.write(line)

    if split_num <= 1:
        # 不分割
        path = out_dir + '/' + color + '.txt'
        char_num = write_txt(char_arr, path)
        if char_num > CHAR_NUM_MAX:
            line = '  字符个数：' + f'{char_num}' + f'       ＜===========字符数大于{CHAR_NUM_MAX}' + '\n---------------------\n\n'
        else:
            line = '  字符个数：' + f'{char_num}' + '\n---------------------\n\n'
        info_file.write(line)
    else:
        # 分割图片
        com_arr = np.copy(char_arr)
        if split_mode == SplitMode.VERTICAL:
            step = math.ceil(height / split_num)
            end = height
        elif split_mode == SplitMode.HORIZONTAL:
            step = math.ceil(width / split_num)
            end = width
        else:
            print("分割时出现了水逆！")
            input("按回车键退出")
            sys.exit(-1)

        i = 0
        part = 0            # 分割图片编号
        while i < end:
            frag_arr = np.zeros((height, width)).astype(dtype=np.str)
            frag_arr[:] = bg_char

            if split_mode == SplitMode.VERTICAL:
                frag_arr[i:i + step, :] = com_arr[i:i + step, :]
            elif split_mode == SplitMode.HORIZONTAL:
                frag_arr[:, i:i + step] = com_arr[:, i:i + step]

            i += step
            part += 1
            part_str = '{:d}'.format(part).zfill(2)
            path = out_dir + '/' + color + '_' + part_str + '.txt'
            char_num = write_txt(frag_arr, path)
            if char_num > CHAR_NUM_MAX:
                line = f'  part{part_str}字符个数：' + f'{char_num}' + f'       ＜===========字符数大于{CHAR_NUM_MAX}' + '\n\n'
            else:
                line = f'  part{part_str}字符个数：' + f'{char_num}' + '\n\n'
            info_file.write(line)

        info_file.write('---------------------\n')


def get_bg_color(r, g, b):
    """
    获取与输入颜色区分度较高的背景颜色
    :param r: 输入颜色R通道
    :param g: 输入颜色G通道
    :param b: 输入颜色B通道
    :return: 背景颜色
    """
    color = np.array([r, g, b])
    mid = np.array([127, 127, 127])
    if np.allclose(color, mid, atol=20, rtol=0):
        # 当前景色比较接近中间灰度时, 修改最大值的通道为255, 增加区分度
        bg_color = np.copy(color)
        bg_color[color.argmax()] = 255
    else:
        bg_color = np.array([255 - r, 255 - g, 255 - b])

    # test
    """
    img = Image.new('RGB', (200, 200), (255, 255, 255))
    arr = np.array(img)
    arr[:, 0:100, :] = color
    arr[:, 100:, :] = bg_color
    img = Image.fromarray(arr)
    img.show()
    """
    return bg_color


def extract_color(img, out_dir):
    """
    提取一张图片的每种颜色, 每种颜色保存为一个或多个(分割)文本文件
    :param img: 要处理的图片
    :param out_dir: 输出路径
    :return: 无
    """
    global info_file

    img_arr = np.array(img)
    img_name = re.split(r'[/ \\]', out_dir)[-1]

    height, width = img_arr.shape[:2]
    if height > H_MAX or width > W_MAX:
        print('---------------------')
        print(f"图片<{img_name}>高为{height}, 宽为{width}, 超出了限制范围(高{H_MAX}x宽{W_MAX}), 这边建议您缩小图片大小重头来过！")
        choice = input("要继续请按 y 并按Enter键, 退出请直接按Enter键：")
        if choice != 'y':
            print("回头是岸, 大侠英明！")
            print('---------------------')
            sys.exit()
        print("你很勇哦！")
        print('---------------------')

    # 保存生成结果的统计信息的文件
    path = out_dir + '/' + 'info.txt'
    info_file = open(path, "w", encoding='UTF-8')
    line = f'图片大小：长={width}, 宽{height}\n---------------------\n'
    info_file.write(line)

    # 判断图片颜色数量, 太大则给出警告提示选择是否结束的信息
    colors = img.getcolors(height * width)
    color_num = len(colors)
    if color_num > COLOR_NUM_MAX:
        print('---------------------')
        print(f"图片<{img_name}>有{color_num}种颜色, 太多了(>{COLOR_NUM_MAX}), 这边建议您减少图片颜色数量重头来过！")
        choice = input("要继续请按 y 并按Enter键, 退出请直接按Enter键：")
        if choice != 'y':
            print("回头是岸, 大侠英明！")
            print('---------------------')
            sys.exit()
        print("你很勇哦！")
        print('---------------------')

    # 遍历图片中包含的颜色进行提取
    for item in colors:
        if not isinstance(item[1], tuple):
            print("图片文件格式不支持！")
            input("按回车键退出")
            continue
        r = item[1][0]
        g = item[1][1]
        b = item[1][2]
        if len(item[1]) == 4:
            a = item[1][3]           # alpha通道, 透明为0, 不透明为255
        else:
            a = -1                   # 无alpha通道

        if a == 0:
            # 完全透明的颜色
            continue

        # RGB(A)合并为一个十六进制数表示
        if a == -1:
            color = '{:02X}'.format(r) + '{:02X}'.format(g) + '{:02X}'.format(b)
            mask = np.where((img_arr[:, :, 0] != r) | (img_arr[:, :, 1] != g) | (img_arr[:, :, 2] != b))
        else:
            color = '{:02X}'.format(r) + '{:02X}'.format(g) + '{:02X}'.format(b) + '{:02X}'.format(a)
            mask = np.where((img_arr[:, :, 0] != r) | (img_arr[:, :, 1] != g) | (img_arr[:, :, 2] != b) |
                            (img_arr[:, :, 3] != a))

        gen_txt(img_arr.shape, color, mask, out_dir)

        # 提取的单色图层预览图
        sep_arr = np.copy(img_arr)
        mask_color = get_bg_color(r, g, b)

        if a != -1:
            mask_color = np.append(mask_color, a)

        sep_arr[mask] = mask_color
        sep_img = Image.fromarray(sep_arr)
        if a == 255:
            color = color[:6]
        sep_img_path = out_dir + '/' + color + '.png'
        sep_img.save(sep_img_path)

    info_file.close()


def proc(in_dir, out_dir="./out"):
    """
    批量处理指定目录中的图片
    :param in_dir: 要处理的图片所在目录
    :param out_dir: 输出目录
    :return: 无
    """
    for entry in os.scandir(in_dir):
        path = in_dir + "/" + entry.name
        print(f"提取: {path}")
        img = Image.open(path, "r")
        if img.mode != 'RGBA':
            img = img.convert('RGB')

        out_sub_dir = out_dir + "/" + entry.name.rsplit(".", 1)[0]
        extract_color(img, out_sub_dir)


def show_settings():
    # 显示设置信息
    global end_mode, split_mode, split_num, fg_char, bg_char

    if end_mode == EndMode.RM_SPACE:
        mode = "消去行尾空格(保留第一行与最后一行行尾空格)"
    elif end_mode == EndMode.KEEP_SPACE:
        mode = "保留行尾空格"

    print(f"  行尾空格处理模式为： {mode}")

    line = "  是否分割图片: "
    if split_num <= 1:
        line += "不分割"
    else:
        if split_mode == SplitMode.VERTICAL:
            mode = "垂直"
        elif split_mode == SplitMode.HORIZONTAL:
            mode = "水平"
        line += f"{mode}分割为 {split_num} 部分"
    print(line)

    print(f"  前景字符: '{fg_char}'")
    print(f"  背景字符: '{bg_char}'")


def get_input():
    # 用户交互
    global end_mode, split_mode, split_num, fg_char, bg_char

    print('---------------------')
    print('当前设置如下：')
    show_settings()
    print('---------------------')

    print("是否采用当前设置：  y.是   n.否")
    choice = input("请选择：")
    print('---------------------')

    modify_flag = False
    if choice == 'n':
        while(True):
            print('---------------------')
            print("请选择要修改的设置：")
            print("  1 修改行尾空格处理模式")
            print("  2 修改图片分割设置")
            print("  3 修改字符")
            print("  0 结束修改")
            choice = input("请选择: ")

            print('---------------------')
            if choice == '0':
                break

            elif choice == '1':
                print("行尾空格处理模式: ")
                print(f"  {EndMode.KEEP_SPACE.value} 保留行尾空格")
                print(f"  {EndMode.RM_SPACE.value} 消去行尾空格(保留第一行与最后一行行尾空格)")
                choice = input("请选择：")
                if end_mode.value != int(choice):
                    modify_flag = True
                if int(choice) == EndMode.RM_SPACE.value:
                    end_mode = EndMode.RM_SPACE
                else:
                    end_mode = EndMode.KEEP_SPACE

            elif choice == '2':
                while True:
                    choice = input("请输入图片分隔份数(0~20, 0或1不做分隔)：")
                    if split_num != int(choice):
                        modify_flag = True
                    if 2 <= int(choice) <= 20:
                        split_num = int(choice)
                        choice = input(f"请输入图片分割模式({SplitMode.VERTICAL.value} 垂直分割, "
                                       f"{SplitMode.HORIZONTAL.value} 水平分割):")
                        if split_mode.value != int(choice):
                            modify_flag = True
                        if int(choice) == SplitMode.VERTICAL.value:
                            split_mode = SplitMode.VERTICAL
                        else:
                            split_mode = SplitMode.HORIZONTAL
                        break
                    elif 0 <= int(choice) <= 1:
                        split_num = int(choice)
                        break
                    else:
                        print("输入超出范围! 请重新输入0~20之间的整数!")

            elif choice == '3':
                while True:
                    print("修改字符设置(注意背景字符和前景字符需要相同宽度)：")
                    print(f"  1 修改前景字符('{fg_char}')")
                    print(f"  2 修改背景字符('{bg_char}')")
                    print(f"  0 结束修改")
                    choice = input("请选择：")
                    if choice == '0':
                        break
                    elif choice == '1':
                        modify_flag = True
                        fg_char = input("请输入前景字符: ")
                    elif choice == '2':
                        modify_flag = True
                        bg_char = input("请输入背景字符: ")

            print('---------------------')
            show_settings()

        print("当前设置如下:")
        show_settings()
        print('---------------------')

    if modify_flag:
        print("是否保存当前设置：")
        print("  1 保存")
        print("  2 不保存")
        choice = input("请选择: ")
        if choice == '1':
            conf = {'end_mode': end_mode.value, 'split_mode': split_mode.value, 'split_num': split_num,
                    'fg_char': fg_char, 'bg_char': bg_char}
            with open(config_path, "w") as f:
                json.dump(conf, f)


def get_config():
    # 读取设置信息
    global end_mode, split_mode, split_num, fg_char, bg_char

    if not os.path.exists(config_path):
        print("没找到配置文件, 采用默认配置")
        return

    with open(config_path, "r") as f:
        conf = json.load(f)
        if int(conf['end_mode']) == EndMode.RM_SPACE.value:
            end_mode = EndMode.RM_SPACE
        elif int(conf['end_mode']) == EndMode.KEEP_SPACE.value:
            end_mode = EndMode.KEEP_SPACE

        if int(conf['split_mode']) == SplitMode.VERTICAL.value:
            split_mode = SplitMode.VERTICAL
        elif int(conf['split_mode']) == SplitMode.HORIZONTAL.value:
            split_mode = SplitMode.HORIZONTAL

        split_num = int(conf['split_num'])
        fg_char = conf['fg_char']
        bg_char = conf['bg_char']


def logo():
    print(r"                       ______                   ")
    print(r"        /\            |  ____|                  ")
    print(r"       /  \      ___  | |__    _   _   _ __     ")
    print(r"      / /\ \    / __| |  __|  | | | | |  _ \    ")
    print(r"     / ____ \  | (__  | |     | |_| | | | | |   ")
    print(r"    /_/    \_\  \___| |_|      \____| |_| |_|   ")
    print(r"                  AC在, 爱一直在!                 ")
    print(r"────────────────────────────────────────────────")
    print(r"AcFun高级弹幕像素画批量颜色转文本工具(v1.2)")
    print(r"────────────────────────────────────────────────")


def show_instruct():
    print("功能简介：")
    print("1. 图片中的每种颜色按像素转换为指定字符组成的文本文件")
    print("2. 将pic文件夹中的图片批量处理, 每张图片的处理结果保存在\n"
          "   out文件夹与图片同名的子文件夹中")
    print("3. 统计生成文本文件中字符个数,便于判断是否超出字符数限制")

    print("\n使用说明：")
    print("1. 在本程序所在目录新建'pic'文件夹")
    print("2. 将要处理的像素图复制到pic文件夹中")
    print("3. 运行本程序, 根据提示进行操作")

    print("\n注意事项：")
    print("1. 目录中不要有其它重要文件以免被覆盖")
    print("2. 再次使用本程序时, 先将已生成的结果备份，以免被覆盖")
    print("3. pic文件夹中只放图片文件, 不要放入其它文件夹或其它文件")
    print("4. 图片颜色不宜太多(建议不超过20种)")
    print("5. 图片的大小不宜过大(建议高和宽不要超100个像素")
    print("6. 分割份数不宜设置过多")
    print(r"────────────────────────────────────────────────")


if __name__ == "__main__":
    in_dir = "./pic"
    out_dir = "./out"

    logo()
    show_instruct()

    get_config()
    get_input()
    init(in_dir, out_dir)

    proc(in_dir, out_dir)

    print("处理完毕")
    input("按回车键退出")
